package com.maple.webrtc.utils;

import java.io.IOException;
import java.io.InputStream;
import java.security.KeyManagementException;
import java.security.KeyStore;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;

import javax.net.ssl.KeyManager;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocketFactory;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;
import javax.net.ssl.X509TrustManager;
public class SSLSocketClient {

    public static class SSLParams {
        public SSLSocketFactory sslSocketFactory;
        public X509TrustManager trustManager;
    }

    public static SSLParams getSSLSocketFactory(
            X509TrustManager trustManager,
            InputStream bksFile,
            String password,
            InputStream... certificates) {
        SSLParams sslParams = new SSLParams();
        try {
            KeyManager[] keyManagers = prepareKeyManager(bksFile, password);
            TrustManager[] trustManagers = prepareTrustManager(certificates);

            // Use user-defined TrustManager if available, otherwise use default or unsafe TrustManager
            TrustManager manager = trustManager != null ? trustManager : getTrustManager(trustManagers);
            if (manager == null) manager = new TrustAllManager();

            // Create TLS SSLContext
            SSLContext sslContext = SSLContext.getInstance("TLS");
            sslContext.init(keyManagers, new TrustManager[]{manager}, null);

            // Get SSLSocketFactory from SSLContext
            sslParams.sslSocketFactory = sslContext.getSocketFactory();
            sslParams.trustManager = (X509TrustManager) manager;
        } catch (NoSuchAlgorithmException | KeyManagementException e) {
            e.printStackTrace();
        }
        return sslParams;
    }

    private static KeyManager[] prepareKeyManager(InputStream bksFile, String password) {
        if (bksFile == null || password == null) return null;
        try {
            KeyStore keyStore = KeyStore.getInstance("BKS");
            keyStore.load(bksFile, password.toCharArray());

            KeyManagerFactory keyManagerFactory = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
            keyManagerFactory.init(keyStore, password.toCharArray());

            return keyManagerFactory.getKeyManagers();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    private static TrustManager[] prepareTrustManager(InputStream... certificates) {
        if (certificates.length == 0) return null;
        try {
            KeyStore keyStore = KeyStore.getInstance(KeyStore.getDefaultType());
            keyStore.load(null);

            CertificateFactory certificateFactory = CertificateFactory.getInstance("X.509");
            int index = 0;

            for (InputStream certificate : certificates) {
                String certAlias = Integer.toString(index++);
                X509Certificate cert = (X509Certificate) certificateFactory.generateCertificate(certificate);
                keyStore.setCertificateEntry(certAlias, cert);
                try {
                    if (certificate != null) certificate.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }

            TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
            trustManagerFactory.init(keyStore);
            return trustManagerFactory.getTrustManagers();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    private static TrustManager getTrustManager(TrustManager[] trustManagers) {
        if (trustManagers == null) return null;
        for (TrustManager trustManager : trustManagers) {
            if (trustManager instanceof X509TrustManager) {
                return (X509TrustManager) trustManager;
            }
        }
        return null;
    }


    public static SSLSocketFactory createSSLSocketFactory() throws Exception {
        SSLContext sc = SSLContext.getInstance("TLS");
        sc.init(null, new TrustManager[]{new TrustAllManager()}, new SecureRandom());
        return sc.getSocketFactory();
    }
}
